package com.onewo.arthasdemo.handler;

import com.onewo.arthasdemo.server.DiscardServerHandler;
import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;

import java.io.ByteArrayInputStream;
import java.lang.reflect.Modifier;
import java.util.HashMap;
import java.util.Map;
import java.util.function.Function;
import java.util.logging.Logger;

public class ClazzOps {

    private static final Logger LOG = Logger.getLogger(ClazzOps.class.getName());

    public static final String OPS_SPEND_TIME = "ops_spend_time";

    public static Map<String, Function<byte[],byte[]>> clazzOpsMap = new HashMap<>();

    private static Function<byte[],byte[]> spentTimeFunction = (classBytes)->{
        CtClass ctClass = null;
        try {
            ctClass = ClassPool.getDefault().makeClass(new ByteArrayInputStream(classBytes));
            for(CtMethod method : ctClass.getMethods()) {
                if (Modifier.isNative(method.getModifiers())) {
                    continue;
                }
                LOG.info("method to be modified:"+ method.getName());
                method.addLocalVariable("cost", CtClass.longType);
                method.insertBefore("cost = System.currentTimeMillis();");
                method.insertAfter(
                        "cost = System.currentTimeMillis() - cost; System.out.println(\"method " + method.getName()
                                + " total cost: \" + cost);");
            }
            return ctClass.toBytecode();
        } catch (Exception e) {
            e.printStackTrace();
            return classBytes;
        }
    };

    static {
        clazzOpsMap.put(OPS_SPEND_TIME,spentTimeFunction);
    }



}
